package yuekao2.machine;

import com.alibaba.alink.operator.batch.BatchOperator;
import com.alibaba.alink.operator.batch.dataproc.SplitBatchOp;
import com.alibaba.alink.operator.batch.evaluation.EvalRegressionBatchOp;
import com.alibaba.alink.operator.batch.sink.CsvSinkBatchOp;
import com.alibaba.alink.operator.batch.source.CsvSourceBatchOp;
import com.alibaba.alink.operator.common.evaluation.RegressionMetrics;
import com.alibaba.alink.operator.common.evaluation.TuningRegressionMetric;
import com.alibaba.alink.pipeline.PipelineModel;
import com.alibaba.alink.pipeline.feature.OneHotEncoder;
import com.alibaba.alink.pipeline.regression.DecisionTreeRegressor;
import com.alibaba.alink.pipeline.regression.LinearRegression;
import com.alibaba.alink.pipeline.tuning.GridSearchCV;
import com.alibaba.alink.pipeline.tuning.ParamGrid;
import com.alibaba.alink.pipeline.tuning.RegressionTuningEvaluator;

public class CarPrice_Assignment {
    public static void main(String[] args) throws Exception {
        //数据准备：调用Alink API加载二手车数据集，提取合适字段（不少二8个）作为特征 features，以及确定目标值target，控制台输出样本数据量；（5分）
        String filePath = "data/yk2/CarPrice_Assignment.csv";
        String schema
                //car_ID,symboling,CarName,fueltype,aspiration,doornumber,carbody,drivewheel,enginelocation,wheelbase,carlength,carwidth
                // ,carheight,curbweight,enginetype,cylindernumber,enginesize,fuelsystem,boreratio,stroke,compressionratio,horsepower,peakrpm,citympg,highwaympg,price
                = "car_ID int,symboling int,CarName String,fueltype String,aspiration String,doornumber String,carbody String,drivewheel String,enginelocation String," +
                "wheelbase double,carlength double,carwidth double,carheight double,curbweight double,enginetype String," +
                "cylindernumber String,enginesize int,fuelsystem String,boreratio double,stroke double,compressionratio double,horsepower double,peakrpm double,citympg double,highwaympg double,price double";

        CsvSourceBatchOp csvSource = new CsvSourceBatchOp()
                .setFilePath(filePath)
                .setSchemaStr(schema)
                .setFieldDelimiter(",");
//                .setLenient(true)
//                .setSkipBlankLine(true)
//                .setIgnoreFirstLine(true);

//        csvSource.print();
        String[] features = new String[]{"aspiration", "doornumber", "carbody", "drivewheel", "carheight", "curbweight", "enginetype", "cylindernumber"};
        String label = "price";

        System.out.println("样本数量:" + csvSource.count());
        //（2）、特征工程：调用Alink 中热编码One-Hot编码，对类别特征进行转换，并且组装为特征向量，划分数据集为训练数据trainData和测试数据集testData；（5分）
        OneHotEncoder one_hot = new OneHotEncoder()
                .setSelectedCols(features)
                .setOutputCols("fec");

        BatchOperator<?> transform = one_hot.fit(csvSource).transform(csvSource);

        BatchOperator<?> spliter = new SplitBatchOp().setFraction(0.8);
        BatchOperator<?> trainData = spliter.linkFrom(transform);
        BatchOperator<?> testData = spliter.getSideOutput(0);
        //（3）、模型训练：分别使用线性回归和决策树回归，设置合理参数值，基于训练数据集构建模型，分别将模型打印控制台，观察参数的值；（5分）
        //线性回归
        LinearRegression lr = new LinearRegression()
                .setVectorCol("fec")
                .setLabelCol(label)
                .setPredictionCol("pred")
                .setMaxIter(3)
                .setNumThreads(1)
                .enableLazyPrintModelInfo();
        BatchOperator<?> lr_transform1 = lr.fit(trainData).transform(testData);
        //决策树回归
        DecisionTreeRegressor dtr = new DecisionTreeRegressor()
                .setPredictionCol("pred")
                .setLabelCol(label)
                .setFeatureCols(features)
                .setMaxBins(128)
                .setMaxDepth(3)
                .enableLazyPrintModelInfo();
        BatchOperator<?> dtr_transform1 = dtr.fit(trainData).transform(testData);


        //（4）、模型评估：合理设置线性回归和决策树回归超参数值，使用trainData训练模型和testData模型预测，并且使用RMSE评估模型，获取最佳模型，并保持文件；（5分）
        RegressionTuningEvaluator regressionTuningEvaluator = new RegressionTuningEvaluator()
                .setPredictionCol("pred")
                .setLabelCol(label)
                .setTuningRegressionMetric(TuningRegressionMetric.RMSE);

        ParamGrid paramGrid1 = new ParamGrid()
                .addGrid(lr, LinearRegression.MAX_ITER, new Integer[]{1, 2, 3})
                .addGrid(lr, LinearRegression.NUM_THREADS, new Integer[]{3, 6, 9});

        GridSearchCV cv1 = new GridSearchCV()
                .setEstimator(lr)
                .setParamGrid(paramGrid1)
                .setTuningEvaluator(regressionTuningEvaluator)
                .setNumFolds(2)
                .enableLazyPrintTrainInfo("TrainInfo");

        ParamGrid paramGrid2 = new ParamGrid()
                .addGrid(dtr, DecisionTreeRegressor.MAX_BINS, new Integer[]{1, 2, 3})
                .addGrid(dtr, DecisionTreeRegressor.MAX_DEPTH, new Integer[]{3, 6, 9});

        GridSearchCV cv2 = new GridSearchCV()
                .setEstimator(dtr)
                .setParamGrid(paramGrid2)
                .setTuningEvaluator(regressionTuningEvaluator)
                .setNumFolds(2)
                .enableLazyPrintTrainInfo("TrainInfo");

        PipelineModel lr_bestPipelineModel = cv1.fit(trainData).getBestPipelineModel();
        BatchOperator<?> lr_transform2 = lr_bestPipelineModel.transform(testData);

        PipelineModel dtr_bestPipelineModel = cv2.fit(trainData).getBestPipelineModel();
        BatchOperator<?> dtr_transform2 = dtr_bestPipelineModel.transform(testData);


        EvalRegressionBatchOp metrics = new EvalRegressionBatchOp()
                .setPredictionCol("pred")
                .setLabelCol(label);

        RegressionMetrics lr_regressionMetrics = metrics.linkFrom(lr_transform2).collectMetrics();
        RegressionMetrics dtr_regressionMetrics = metrics.linkFrom(dtr_transform2).collectMetrics();

        System.out.println("lr RMSE"+lr_regressionMetrics.getMse());
        System.out.println("dtr RMSE"+dtr_regressionMetrics.getMse());

        if (lr_regressionMetrics.getMse()>dtr_regressionMetrics.getMse()){
            System.out.println("线性回归较好");
            lr_bestPipelineModel.save("data/yk2/lr.csv");
            lr_transform2.link(new CsvSinkBatchOp()
                    .setFilePath("data/yk2/lr_test.txt"));
        }else{
            System.out.println("决策树回归较好");
            dtr_bestPipelineModel.save("data/yk2/dtr.csv");
            dtr_transform2.link(new CsvSinkBatchOp()
                    .setFilePath("data/yk2/dtr_test.txt"));
        }
    }
}
